package com.huiying.config;

import com.huiying.constant.JwtClaimsConstant;
import com.huiying.context.BaseContext;
import com.huiying.properties.JwtProperties;
import com.huiying.utils.JwtUtil;
import io.jsonwebtoken.Claims;
import io.jsonwebtoken.ExpiredJwtException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.config.ChannelRegistration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;

/**
 * @Author: lsh
 * @Date: 2024/12/08/13:57
 * @Description: WebSocket配置类
 */
@Configuration
@EnableWebSocketMessageBroker // 启用 WebSocket 消息代理
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {

    private static final Logger log = LoggerFactory.getLogger(WebSocketConfig.class);

    @Autowired // 注入 JWT 属性配置类
    private JwtProperties jwtProperties;

    @Autowired // 注入 RedisTemplate 用于验证 Redis 中的令牌
    private RedisTemplate<String, Object> redisTemplate; // 注意这里泛型可以根据您实际存储类型调整

    @Override
    public void configureMessageBroker(MessageBrokerRegistry config) {
        config.enableSimpleBroker("/queue", "/topic"); // 消息代理前缀
        config.setApplicationDestinationPrefixes("/app"); // 应用目的地前缀
        config.setUserDestinationPrefix("/user"); // 用户目的地前缀，必须配置以支持 /user/{userId}/queue/messages 这种形式
    }

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint("/ws-chat").setAllowedOriginPatterns(
                "http://localhost:*",       // 允许所有本地端口，如 http://localhost:3000, http://localhost:8081
                "null",                     // 允许本地文件直接访问 (file://协议)
                "https://servicewechat.com" // 微信小程序可能需要的源
        );
    }

    @Override
    public void configureClientInboundChannel(ChannelRegistration registration) {
        registration.interceptors(new ChannelInterceptor() {
            /**
             * 在消息发送到 Channel 之前调用。
             * 这是处理 STOMP CONNECT 帧并设置用户 Principal 的地方。
             * @param message 待发送的消息
             * @param channel 消息将要发送到的 Channel
             * @return 消息本身（继续处理）或 null（拒绝处理）
             */
            @Override
            public Message<?> preSend(Message<?> message, MessageChannel channel) {
                StompHeaderAccessor accessor =
                        MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);

                // 只处理 STOMP CONNECT 帧进行认证
                if (StompCommand.CONNECT.equals(accessor.getCommand())) {
                    // 从 STOMP CONNECT 帧的头部获取 Authorization 令牌
                    // 客户端需要在 stompClient.connect() 的第一个参数中传递 {'Authorization': 'Bearer YOUR_JWT_TOKEN'}
                    String token = accessor.getFirstNativeHeader("Authorization");

                    if (token != null && token.startsWith("Bearer ")) {
                        token = token.substring(7); // 移除 "Bearer " 前缀

                        try {
                            // 1. 解析 JWT 令牌
                            Claims claims = JwtUtil.parseJWT(jwtProperties.getUserSecretKey(), token);
                            Long userId = Long.valueOf(claims.get(JwtClaimsConstant.USER_ID).toString());

                            // 2. 从 Redis 验证令牌是否有效（防止多点登录或过期）
                            // 这里的 Redis key 应该与您 JwtTokenUserInterceptor 中生成/存储的规则一致
                            String redisKey = "user:token:" + userId; // 假设您的 Redis key 格式是 "user:token:userId"
                            String redisToken = (String) redisTemplate.opsForValue().get(redisKey);

                            if (redisToken != null && redisToken.equals(token)) {
                                // 认证成功，设置用户 Principal
                                // 这个 Principal 会在后续的 @MessageMapping 方法中通过 SimpMessageHeaderAccessor.getUser() 获取
                                accessor.setUser(new UsernamePasswordAuthenticationToken(userId.toString(), null, null));
                                log.info("WebSocket用户 {} 已通过CONNECT帧JWT认证并设置为主体。", userId);
                            } else {
                                // Redis 中没有令牌或令牌不匹配，表示无效、已过期或在其他地方登录
                                log.warn("WebSocket连接尝试被拒绝：JWT令牌无效或已过期（Redis检查失败）。Token: {}", token);
                                return null; // 拒绝连接
                            }

                        } catch (ExpiredJwtException ex) {
                            log.warn("WebSocket连接尝试被拒绝：JWT令牌已过期。Token: {}, 错误: {}", token, ex.getMessage());
                            return null; // JWT 过期，拒绝连接
                        } catch (Exception ex) {
                            // 处理 JWT 解析异常（如签名无效、格式错误等）
                            log.error("WebSocket连接尝试被拒绝：JWT解析失败。Token: {}, 错误: {}", token, ex.getMessage(), ex);
                            return null; // 其他 JWT 解析错误，拒绝连接
                        }
                    } else {
                        log.warn("WebSocket连接尝试被拒绝：CONNECT帧中未找到有效的JWT令牌或格式不正确。");
                        return null; // 未提供 JWT 或格式不正确，拒绝连接
                    }
                }
                // 对于 STOMP DISCONNECT 帧，不需要 BaseContext 相关的清理了
                // 因为我们不再依赖 BaseContext 来存储 Principal
                else if (StompCommand.DISCONNECT.equals(accessor.getCommand())) {
                    // 当客户端发送 DISCONNECT 帧时，记录日志
                    // 用户信息可以从 accessor.getUser() 获取
                    log.info("WebSocket连接断开命令收到，用户ID: {}", accessor.getUser() != null ? accessor.getUser().getName() : "未知");
                }

                // 继续处理其他 STOMP 消息（SEND, SUBSCRIBE 等）
                // 它们的认证信息会从 accessor.getUser() 中获取，不需要再次处理
                return message;
            }

            /**
             * 在消息发送完成后调用，无论发送成功与否。
             * 这是清理 ThreadLocal 的重要时机，以避免内存泄漏和数据污染。
             * 对于 WebSocket 连接，主要用于清理与 Principal 相关的上下文。
             * @param message 已经发送的消息
             * @param channel 消息被发送到的 Channel
             * @param sent 是否成功发送
             * @param ex 如果发送失败，发生的异常
             */
            @Override
            public void afterSendCompletion(Message<?> message, MessageChannel channel, boolean sent, Exception ex) {
                StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
                // 确保 Principal 在消息处理链结束后被清除
                if (accessor != null && accessor.getUser() != null) {
                    log.debug("WebSocket消息处理完成，用户 Principal 清理完毕: {}", accessor.getUser().getName());
                    // 如果您有其他与 WebSocket 消息处理线程绑定的 ThreadLocal，可以在这里进行清理
                }
            }
        });
    }
}